import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.preprocessing.image import ImageDataGenerator

# 设置路径
train_path = 'dataset/train'
validation_path = 'dataset/test'

# 数据预处理和增强
train_datagen = ImageDataGenerator(
    rescale=1. / 255,
    shear_range=0.2,
    zoom_range=0.2,
    horizontal_flip=True
)

validation_datagen = ImageDataGenerator(rescale=1. / 255)

train_generator = train_datagen.flow_from_directory(
    train_path,
    target_size=(224, 224),  # 调整为模型所需的输入尺寸
    batch_size=32,
    class_mode='categorical'
)

validation_generator = validation_datagen.flow_from_directory(
    validation_path,
    target_size=(224, 224),
    batch_size=32,
    class_mode='categorical'
)

# 定义残差模型
input_shape = (224, 224, 3)
num_classes = len(train_generator.class_indices)


def residual_block(x, filters, downsample=False):
    strides = (2, 2) if downsample else (1, 1)

    # 主分支
    y = layers.Conv2D(filters, kernel_size=(3, 3), strides=strides, padding='same')(x)
    y = layers.BatchNormalization()(y)
    y = layers.Activation('relu')(y)

    y = layers.Conv2D(filters, kernel_size=(3, 3), padding='same')(y)
    y = layers.BatchNormalization()(y)

    # 跳跃连接分支
    if downsample:
        x = layers.Conv2D(filters, kernel_size=(1, 1), strides=strides, padding='same')(x)
        x = layers.BatchNormalization()(x)

    # 残差相加
    out = layers.Add()([x, y])
    out = layers.Activation('relu')(out)
    return out


def build_resnet():
    inputs = keras.Input(shape=input_shape)
    x = layers.Conv2D(64, kernel_size=(7, 7), strides=(2, 2), padding='same')(inputs)
    x = layers.BatchNormalization()(x)
    x = layers.Activation('relu')(x)
    x = layers.MaxPooling2D(pool_size=(3, 3), strides=(2, 2), padding='same')(x)

    x = residual_block(x, filters=64)
    x = residual_block(x, filters=64)
    x = residual_block(x, filters=64)

    x = layers.GlobalAveragePooling2D()(x)
    x = layers.Dense(256, activation='relu')(x)
    outputs = layers.Dense(num_classes, activation='softmax')(x)

    model = keras.Model(inputs=inputs, outputs=outputs)
    return model


# 构建并编译残差模型
model = build_resnet()
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

# 训练模型
history = model.fit(
    train_generator,
    steps_per_epoch=len(train_generator),
    epochs=10,
    validation_data=validation_generator,
    validation_steps=len(validation_generator)
)

# 在测试集上评估模型
test_loss, test_accuracy = model.evaluate(validation_generator, steps=len(validation_generator))
print("Test Loss:", test_loss)
print("Test Accuracy:", test_accuracy)

import matplotlib.pyplot as plt

# 绘制损失函数图像
plt.plot(history.history['loss'], label='Train Loss')
plt.plot(history.history['val_loss'], label='Validation Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.show()

# 绘制准确率图像
plt.plot(history.history['accuracy'], label='Train Accuracy')
plt.plot(history.history['val_accuracy'], label='Validation Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()
plt.show()

import numpy as np
from sklearn.metrics import classification_report

# 在测试集上进行预测
y_pred = model.predict(validation_generator)
y_pred = np.argmax(y_pred, axis=1)  # 将预测概率转换为类别索引
y_true = validation_generator.classes

# 计算评价指标
report = classification_report(y_true, y_pred, target_names=validation_generator.class_indices.keys())

# 打印评价指标
print(report)